#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
from scene import Scene
import os
import numpy as np
from random import randint
from tqdm import tqdm
from os import makedirs
from gaussian_renderer import render
import torchvision
from utils.general_utils import safe_state
from argparse import ArgumentParser
from arguments import ModelParams, PipelineParams, get_combined_args
from gaussian_renderer import GaussianModel
import copy, time, imageio

normalize = lambda x: (x - x.min()) / (x.max() - x.min())
to8b = lambda x: x.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()


def generate_multi_views(view_source, length=60):
    results = []
    angle0, angle1 = 0, 60
    theta0 = angle0 * np.pi / 180
    theta1 = angle1 * np.pi / 180
    for i in range(length):
        view = copy.deepcopy(view_source)
        
        weight = (i+1) / length
        theta = theta0 * (1-weight) + theta1 * weight
        sin_theat = np.sin(theta)
        cos_theat = np.cos(theta)
        R_offset = torch.tensor([[1,    0,          0],
                                 [0, cos_theat, -sin_theat],
                                 [0, sin_theat, cos_theat]], 
                                dtype=view.world_view_transform.dtype, 
                                device=view.world_view_transform.device)
        
        R_hat = view.world_view_transform[:3, :3] @ R_offset
        view.world_view_transform[:3, :3] = R_hat
        
        results.append(view)
    return results

def render_vedio(model_path, train_views, gaussians, pipeline, background):
    scene_name = model_path.split('/')[-1]
    idx = randint(0, len(train_views)-1)
    views = train_views[idx]
    
    length_view = 180
    generated_views = generate_multi_views(views, length_view)
    
    video_images_path = os.path.join(model_path, f"test_images") #
    makedirs(video_images_path, exist_ok=True)
    
    render_video_out = imageio.get_writer(f'{video_images_path}/{scene_name}_test.mp4', mode='I', fps=30,codec='libx264',quality=10.0)#
    
    for idx, view in enumerate(tqdm(generated_views, desc="Rendering progress")):
                
        rendered_pkg = render(view, gaussians, pipeline, background)
        rendering, depth, weight = rendered_pkg['render'], rendered_pkg['depth'], rendered_pkg['weight']
        np_image = to8b(rendering)
        np_depth = to8b(normalize(depth)[None].repeat(3,1,1))
        np_weight = to8b(normalize(weight)[None].repeat(3,1,1))
        
        render_video_out.append_data(np.concatenate([np_image, np_depth, np_weight],axis=1))
        
    render_video_out.close()

def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, render_multiview_vedio: bool):
    with torch.no_grad():
        gaussians = GaussianModel(dataset.sh_degree)
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)

        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        model_path = dataset.model_path

        if not render_multiview_vedio:
            render_vedio(model_path, scene.getTrainCameras(), gaussians, pipeline, background)

if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser(description="Testing script parameters")
    model = ModelParams(parser, sentinel=True)
    pipeline = PipelineParams(parser)
    parser.add_argument("--iteration", default=-1, type=int)
    parser.add_argument("--skip_train", action="store_true")
    parser.add_argument("--skip_test", action="store_true")
    parser.add_argument("--render_multiview_vedio", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    args = get_combined_args(parser)
    print("Rendering " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    render_sets(model.extract(args), 
                args.iteration, 
                pipeline.extract(args), 
                args.render_multiview_vedio)